#include <iostream>
#include <winsock2.h>
#include <windows.h>
#include <mmsystem.h>
#include <vector>

#pragma comment(lib, "ws2_32.lib")
#pragma comment(lib, "winmm.lib")

#define RTP_PORT 5004
#define RTP_PACKET_SIZE 1024
#define BUFFER_COUNT 4 // 使用4个缓冲区来更好地处理抖动
#define PREBUFFER_PACKETS 50 // 预缓存 50 个包（约 1 秒）

int main() {
    WSADATA wsaData;
    SOCKET sock;
    sockaddr_in addr, sender;
    int senderLen = sizeof(sender);
    WAVEFORMATEX wfx = {0};
    HWAVEOUT hWaveOut = nullptr;

    // 初始化 Winsock
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        std::cerr << "WSAStartup failed" << std::endl;
        return 1;
    }

    // 设置音频格式（与发送端一致）
    wfx.wFormatTag = WAVE_FORMAT_PCM;
    wfx.nChannels = 1;
    wfx.nSamplesPerSec = 44100;
    wfx.wBitsPerSample = 16;
    wfx.nBlockAlign = wfx.nChannels * wfx.wBitsPerSample / 8;
    wfx.nAvgBytesPerSec = wfx.nSamplesPerSec * wfx.nBlockAlign;
    wfx.cbSize = 0;

    // 创建 UDP socket
    sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock == INVALID_SOCKET) {
        std::cerr << "socket 创建失败" << std::endl;
        return 1;
    }
    addr.sin_family = AF_INET;
    addr.sin_port = htons(RTP_PORT);
    addr.sin_addr.s_addr = INADDR_ANY;
    if (bind(sock, (sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR) {
        std::cerr << "bind 失败" << std::endl;
        closesocket(sock);
        WSACleanup();
        return 1;
    }

    // 打开音频输出设备
    if (waveOutOpen(&hWaveOut, WAVE_MAPPER, &wfx, 0, 0, CALLBACK_NULL) != MMSYSERR_NOERROR) {
        std::cerr << "无法打开音频输出设备" << std::endl;
        closesocket(sock);
        WSACleanup();
        return 1;
    }
    std::cout << "音频设备已打开，等待接收数据..." << std::endl;

    // 创建并准备多个缓冲区
    std::vector<WAVEHDR> waveHeaders(BUFFER_COUNT);
    std::vector<std::vector<BYTE>> buffers(BUFFER_COUNT, std::vector<BYTE>(RTP_PACKET_SIZE));
    for (int i = 0; i < BUFFER_COUNT; ++i) {
        waveHeaders[i] = {0};
        waveHeaders[i].lpData = reinterpret_cast<LPSTR>(buffers[i].data());
        waveHeaders[i].dwBufferLength = RTP_PACKET_SIZE;
        waveHeaders[i].dwFlags = WHDR_DONE; // 初始状态为完成，可以被使用
        waveOutPrepareHeader(hWaveOut, &waveHeaders[i], sizeof(WAVEHDR));
    }

    // 预缓存
    BYTE preBuffer[PREBUFFER_PACKETS * RTP_PACKET_SIZE];
    int preBuffered = 0;
    while (preBuffered < PREBUFFER_PACKETS) {
        int ret = recvfrom(sock, (char*)(preBuffer + preBuffered * RTP_PACKET_SIZE), RTP_PACKET_SIZE, 0, (sockaddr*)&sender, &senderLen);
        if (ret <= 0) break;
        preBuffered++;
    }
    std::cout << "预缓存完成，开始播放..." << std::endl;
    // 播放预缓存数据
    WAVEHDR waveHeader = {0};
    waveHeader.lpData = reinterpret_cast<LPSTR>(preBuffer);
    waveHeader.dwBufferLength = preBuffered * RTP_PACKET_SIZE;
    waveHeader.dwFlags = 0;
    waveHeader.dwLoops = 0;
    waveHeader.lpNext = nullptr;
    waveHeader.reserved = 0;
    if (waveOutPrepareHeader(hWaveOut, &waveHeader, sizeof(WAVEHDR)) == MMSYSERR_NOERROR) {
        waveOutWrite(hWaveOut, &waveHeader, sizeof(WAVEHDR));
    }
    // 边接收边播放
    BYTE buffer[RTP_PACKET_SIZE];
    while (true) {
        int ret = recvfrom(sock, (char*)buffer, RTP_PACKET_SIZE, 0, (sockaddr*)&sender, &senderLen);
        if (ret <= 0) break;
        WAVEHDR waveHeader2 = {0};
        waveHeader2.lpData = reinterpret_cast<LPSTR>(buffer);
        waveHeader2.dwBufferLength = ret;
        waveHeader2.dwFlags = 0;
        waveHeader2.dwLoops = 0;
        waveHeader2.lpNext = nullptr;
        waveHeader2.reserved = 0;
        if (waveOutPrepareHeader(hWaveOut, &waveHeader2, sizeof(WAVEHDR)) == MMSYSERR_NOERROR) {
            waveOutWrite(hWaveOut, &waveHeader2, sizeof(WAVEHDR));
        }
    }

cleanup:
    std::cout << "音频流结束，清理资源..." << std::endl;
    // 等待所有缓冲区播放完成
    while (true) {
        bool allDone = true;
        for (int i = 0; i < BUFFER_COUNT; ++i) {
            if (!(waveHeaders[i].dwFlags & WHDR_DONE)) {
                allDone = false;
                break;
            }
        }
        if (allDone) {
            break;
        }
        Sleep(10);
    }

    // 取消准备并关闭设备
    for (int i = 0; i < BUFFER_COUNT; ++i) {
        waveOutUnprepareHeader(hWaveOut, &waveHeaders[i], sizeof(WAVEHDR));
    }
    waveOutClose(hWaveOut);
    closesocket(sock);
    WSACleanup();
    std::cout << "播放完成" << std::endl;
    return 0;
}
